"""
HDF5数据处理模块
负责协调HDF5文件处理和绘图操作
"""

import os
import re
import numpy as np
import matplotlib.pyplot as plt
from utils.file_handler import HDF5FileHandler
from core.plot_manager import PlotManager


class DataProcessor:
    """HDF5数据处理类，负责协调文件处理和绘图"""
    
    def __init__(self):
        self.file_handler = HDF5FileHandler()
        self.plot_manager = PlotManager()
    
    def process_file_loaded(self, file_path):
        """
        处理文件加载事件
        
        Args:
            file_path: HDF5文件路径
            
        Returns:
            dict: 处理结果，包含文件结构和采样频率等信息
        """
        try:
            success = self.file_handler.open_file(file_path)
            if not success:
                return {
                    'success': False,
                    'message': '无法打开HDF5文件'
                }
            
            # 获取文件结构
            file_structure = self.file_handler.get_file_structure()
            
            # 获取采样频率
            sampling_freq = self.file_handler.get_sampling_rate()
            
            # 自动查找幅度-相位配对
            amp_phase_pairs = self.file_handler.find_amp_phase_pairs()
            
            return {
                'success': True,
                'file_structure': file_structure,
                'sampling_freq': sampling_freq,
                'amp_phase_pairs': amp_phase_pairs,
                'message': f'成功加载文件，找到 {len(amp_phase_pairs)} 个数据对'
            }
            
        except Exception as e:
            return {
                'success': False,
                'message': f'处理文件时出错: {str(e)}'
            }
    
    def process_sorting_request(self):
        """
        处理数据对自动排序请求
        
        Returns:
            dict: 处理结果，包含排序后的数据对列表
        """
        try:
            # 从 file_handler 获取当前的数据对
            pairs = self.file_handler.find_amp_phase_pairs()
            
            # 定义自然排序的 key function
            def natural_sort_key(s):
                return [int(text) if text.isdigit() else text.lower() 
                        for text in re.split('([0-9]+)', s['base_name'])]

            # 按 base_name 进行自然排序
            sorted_pairs = sorted(pairs, key=natural_sort_key)
            
            return {
                'success': True,
                'sorted_pairs': sorted_pairs,
                'message': f'成功对 {len(sorted_pairs)} 个数据对进行自然排序'
            }
            
        except Exception as e:
            return {
                'success': False,
                'message': f'排序时出错: {str(e)}'
            }
    
    def process_preview_request(self, file_path):
        """
        处理预览请求
        
        Args:
            file_path: 文件路径
            
        Returns:
            str: 预览文本
        """
        try:
            if not self.file_handler.current_file:
                # 如果文件没有打开，尝试打开
                success = self.file_handler.open_file(file_path)
                if not success:
                    return "无法打开文件进行预览"
            
            return self.file_handler.create_preview_text(file_path)
            
        except Exception as e:
            return f"预览文件时出错: {str(e)}"
    
    def process_plot_request(self, parameters):
        """
        处理绘图请求，支持单一波形或批量处理
        
        Args:
            parameters: 包含所有绘图相关参数的字典
            
        Returns:
            dict: 处理结果
        """
        try:
            batch_enabled = parameters.get('batch_enabled', False)
            
            if batch_enabled:
                start_index = parameters.get('start_index', 0)
                end_index = parameters.get('end_index', 0)
                indices = range(start_index, end_index + 1)
                if len(indices) > 100:  # 安全限制，防止处理过多图形
                    return {'success': False, 'message': '批量处理数量过多 (最大100)，请缩小范围。'}
            else:
                indices = [parameters.get('waveform_index', 0)]
                
            if not indices:
                return {'success': False, 'message': '未选择任何波形索引。'}

            # 统一处理流程
            results = []
            for index in indices:
                result = self._plot_single_waveform(parameters, index)
                results.append(result)
                # 在保存模式下，遇到第一个错误就停止
                if not result['success'] and parameters.get('plot_or_save') == 'Save':
                    return {'success': False, 'message': f'波形 {index} 处理失败: {result["message"]}'}
            
            # 总结结果
            success_count = sum(1 for r in results if r['success'])
            if success_count == 0:
                return {'success': False, 'message': f'所有波形处理失败。\n第一个错误: {results[0].get("message", "未知错误")}'}

            if parameters.get('plot_or_save') == 'Plot':
                return {'success': True, 'message': f'已成功提交 {success_count} 个波形图的绘制请求。'}
            else:  # Save
                first_save_path = os.path.dirname(results[0].get('save_path', '')) if success_count > 0 else '未知路径'
                return {'success': True, 'message': f'成功保存 {success_count} 个图形到:\n{first_save_path}'}

        except Exception as e:
            return {
                'success': False,
                'message': f'绘图处理出错: {str(e)}'
            }

    def _plot_single_waveform(self, parameters, waveform_index):
        """
        处理单个波形的绘图或保存
        
        Args:
            parameters: 包含所有绘图相关参数的字典
            waveform_index: 当前要处理的波形索引
            
        Returns:
            dict: 单个处理任务的结果
        """
        fig = None
        try:
            file_path = parameters['file_path']
            selected_pairs = parameters['selected_pairs']
            plot_or_save = parameters['plot_or_save']
            
            if not selected_pairs:
                return {'success': False, 'message': '没有选择要绘制的数据对'}
            
            # 确保文件已打开
            if not self.file_handler.current_file or self.file_handler.current_file.filename != file_path:
                if not self.file_handler.open_file(file_path):
                    return {'success': False, 'message': '无法重新打开HDF5文件'}

            # 准备数据
            data_pairs = []
            sampling_freq = float(parameters.get('sampling_freq', 1.0))
            
            for pair in selected_pairs:
                amp_path = pair['amp_dataset']['path']
                phase_path = pair['phase_dataset']['path']
                
                amp_data = self.file_handler.get_dataset_data(amp_path, waveform_index)
                phase_data = self.file_handler.get_dataset_data(phase_path, waveform_index)
                
                if amp_data is None or phase_data is None:
                    continue
                
                t = np.arange(len(amp_data)) / (sampling_freq * 1e6)
                
                data_pairs.append({
                    'name': pair['base_name'], 'amp_data': amp_data,
                    'phase_data': phase_data, 't': t
                })
            
            if not data_pairs:
                return {'success': False, 'message': '无法读取选中的数据'}

            # 创建图形
            plot_params = parameters.copy()
            plot_params['sampling_freq'] = sampling_freq
            npR, npC = self.plot_manager.get_plot_grid_size(len(data_pairs))
            fig, axes = plt.subplots(npR, npC, figsize=parameters['figure_size'], squeeze=False)
            
            # 绘制内容和标题
            self.plot_manager.plot_figure_content(axes, data_pairs, plot_params)
            file_name = os.path.basename(file_path)
            title = self.plot_manager.create_figure_title(file_name, waveform_index, sampling_freq)
            self.plot_manager.setup_figure_layout(fig, title)
            
            if plot_or_save == 'Plot':
                plt.show()  # 阻塞式，逐个显示
                return {'success': True}
            else:  # Save
                save_path = self._generate_save_path(file_path, waveform_index)
                fig.savefig(save_path, dpi=300, bbox_inches='tight')
                plt.close(fig)
                return {'success': True, 'save_path': save_path}
                    
        except Exception as e:
            if fig:
                plt.close(fig)
            return {'success': False, 'message': str(e)}
    
    def _generate_save_path(self, file_path, waveform_index):
        """
        生成保存路径
        
        Args:
            file_path: 原始文件路径
            waveform_index: 波形索引
            
        Returns:
            str: 保存路径
        """
        dir_path = os.path.dirname(file_path)
        base_name = os.path.splitext(os.path.basename(file_path))[0]
        
        # 创建输出目录
        output_dir = os.path.join(dir_path, f"{base_name}_plots")
        os.makedirs(output_dir, exist_ok=True)
        
        # 生成文件名
        save_name = f"{base_name}_waveform_{waveform_index:03d}.png"
        return os.path.join(output_dir, save_name)
    
    def get_amp_phase_pairs(self):
        """
        获取当前文件的幅度-相位配对
        
        Returns:
            list: 配对列表
        """
        if self.file_handler.current_file:
            return self.file_handler.find_amp_phase_pairs()
        return []
    
    def close_file(self):
        """关闭当前文件"""
        self.file_handler.close_file()
    
    def __del__(self):
        """析构函数，确保文件被正确关闭"""
        self.close_file() 